﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Linq.Expressions;
using System.Reflection;

namespace CSharp4ExpressionMetaCompiler
{
    public class BlockCompilerVisitor : System.Linq.Expressions.ExpressionVisitor
    {
        private Dictionary<ParameterExpression, Dictionary<string, ParameterExpression>> blockStartParameters =
            new Dictionary<ParameterExpression, Dictionary<string, ParameterExpression>>();
        private Dictionary<string, LabelTarget> labelTargets = new Dictionary<string, LabelTarget>();
        private Dictionary<ParameterExpression, ParameterExpression> parametersToBeReplaced = new Dictionary<ParameterExpression, ParameterExpression>();

        public Expression StartVisit(Expression exp)
        {
            return Visit(exp);
        }
        
        protected override Expression VisitParameter(ParameterExpression node)
        {
            ParameterExpression result = null;
            if (parametersToBeReplaced.TryGetValue(node, out result))
                return result;
            return base.VisitParameter(node);
        }

        protected override Expression VisitMember(MemberExpression node)
        {
            if ((node.Member.MemberType == MemberTypes.Property) && (node.Expression is ParameterExpression))
            {
                Dictionary<string, ParameterExpression> parameters = null;
                if (blockStartParameters.TryGetValue(node.Expression as ParameterExpression, out parameters))
                    return parameters[node.Member.Name];
            }
            return base.VisitMember(node);
        }

        protected override Expression VisitMethodCall(MethodCallExpression node)
        {
            if (IsMethodOf<Block>(node))
            {
                var expressions = new List<Expression>();
                do 
                {
                    var r = VisitBlockMethodCall(node as MethodCallExpression);
                    expressions.Insert(0, r);
                    node = node.Object as MethodCallExpression;
                } while (node != null);

                return Expression.Block(expressions);
            }
            return base.VisitMethodCall(node);
        }

        private Expression VisitBlockMethodCall(MethodCallExpression node)
        {
            if (IsMethodOf<Block>(node, "Assign"))
            {
                var left = node.Arguments[0] as LambdaExpression;
                return Expression.Assign(Visit(left.Body), Visit(node.Arguments[1]));
            }
            if (IsMethodOf<Block>(node, "Label"))
            {
                string labelName = (string)(node.Arguments[0] as ConstantExpression).Value;
                return Expression.Label(GetLabelTarget(labelName));
            }
            if (IsMethodOf<Block>(node, "IfThen"))
                return Expression.IfThen(Visit(node.Arguments[0]), Visit((node.Arguments[1] as LambdaExpression).Body));
            if (IsMethodOf<Block>(node, "IfThenElse"))
                return Expression.IfThenElse(Visit(node.Arguments[0]), 
                    Visit((node.Arguments[1] as LambdaExpression).Body),
                    Visit((node.Arguments[2] as LambdaExpression).Body));
            if (IsMethodOf<Block>(node, "Goto"))
            {
                string labelName = (string)(node.Arguments[0] as ConstantExpression).Value;
                return Expression.Goto(GetLabelTarget(labelName));
            }
            if (IsMethodOf<Block>(node, "Loop"))
            {
                var result = (node.Arguments[0] as LambdaExpression).Body;
                return Expression.Loop(Visit(result));
            }
            if (IsMethodOf<Block>(node, "While"))
            {
                var exitLabel = Expression.Label();

                var block =
                    Expression.Block(
                        new Expression[] 
                        {
                            Expression.Loop(
                                Expression.Block(
                                    new Expression[] 
                                    {
                                        Expression.IfThen(
                                            Expression.Not(Visit(node.Arguments[0])),
                                            Expression.Goto(exitLabel)),
                                        Visit((node.Arguments[1] as LambdaExpression).Body)
                                    })),
                            Expression.Label(exitLabel)
                        });
                return block;
            }
            if (IsMethodOf<Block>(node, "_"))
                return Visit((node.Arguments[0] as LambdaExpression).Body);
            if (IsMethodOf<Block>(node, "For"))
            {
                var l = node.Arguments[2] as LambdaExpression;
                var i = l.Parameters[0];
                var test = node.Arguments[1] as LambdaExpression;
                var exitLabel = Expression.Label();

                var block =
                    Expression.Block(new ParameterExpression[] { i },
                        new Expression[]
                        {
                            Expression.Assign(i, Visit(node.Arguments[0])),
                            Expression.Loop(
                                Expression.Block(
                                    new Expression[] 
                                    {
                                        Expression.IfThen(VisitWithReplaceParameter(test.Parameters[0], i, Expression.Not(test.Body)), 
                                            Expression.Goto(exitLabel)),
                                        VisitWithReplaceParameter(l.Parameters[0], i, l.Body),
                                        Expression.Assign(i, Expression.Add(i, Expression.Constant(1)))
                                    })),
                            Expression.Label(exitLabel)
                        });
                return block;
            }
            if (IsMethodOf<Block>(node, "Start", true))
            {
                var newExpression = node.Arguments[0] as NewExpression;
                var variablesToAdd =
                    newExpression.Type.GetProperties()
                        .Select((p, i) =>
                            new
                            {
                                Parameter = Expression.Parameter(p.PropertyType, p.Name),
                                Value = newExpression.Arguments[i]
                            }).ToArray();

                var blockStartLambda = node.Arguments[1] as LambdaExpression;
                var parameters = variablesToAdd.ToDictionary(v => v.Parameter.Name, v => v.Parameter);
                
                blockStartParameters.Add(blockStartLambda.Parameters[0], parameters);
                var body = Visit(blockStartLambda.Body);
                blockStartParameters.Remove(blockStartLambda.Parameters[0]);

                var blockWithVariablesInitialization =
                    variablesToAdd.Select(v => Expression.Assign(v.Parameter, v.Value))
                        .Concat(new Expression[] { body });

                var block = Expression.Block(variablesToAdd.Select(v => v.Parameter), blockWithVariablesInitialization);
                return block;
            }
            if (IsMethodOf<Block>(node, "Declare", true))
            {
                return node.Arguments[1];
            }

            return null; 
        }

        private Expression VisitWithReplaceParameter(ParameterExpression parameterToReplace, ParameterExpression newParameter, Expression expression)
        {
            parametersToBeReplaced.Add(parameterToReplace, newParameter);
            var result = Visit(expression);
            parametersToBeReplaced.Remove(parameterToReplace);
            
            return result;
        }
        private LabelTarget GetLabelTarget(string labelName)
        {
            LabelTarget target = null;
            if (!(labelTargets.TryGetValue(labelName, out target)))
            {
                target = Expression.Label(labelName);
                labelTargets.Add(labelName, target);
            }
            return target;
        }
        private bool IsMethodOf<T>(MethodCallExpression node, string methodName = null, bool isStatic = false)
        {
            if (node == null)
                return false;
            if (methodName == null)
                return node.Method.DeclaringType == typeof(T);
            return ((node.Method.DeclaringType == typeof(T)) && (node.Method.IsStatic == isStatic)
                && (node.Method.Name == methodName));
        }
    }
}
